import sys
sys.stdout.reconfigure(encoding='utf-8')

from pinecone import Pinecone, ServerlessSpec

pinecone = Pinecone(api_key="pcsk_6aQeZm_5UcBB6iZSR4j5gnmoa1UZm2oJdmWuQoTUQqkDiwN7nApxibHDX2VC6Nb1yp2zZP")
# 索引名称
index_name = "mnist-index"

################### 检查、清除之前创建的索引 ###################
# 获取现有索引列表
existing_indexes = pinecone.list_indexes()

# 检查索引是否存在，如果存在就删除
# 这个if是否需要，要看情况而定
# 比如有的时候，如果不想要重复删除在创建，这个if就可以不要
if any(index['name'] == index_name for index in existing_indexes):
    print(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    print(f"索引 '{index_name}' 已成功删除。")
else:
    print(f"索引 '{index_name}' 不存在，将创建新索引。")
################### 检查、清除之前创建的索引 ###################

index_name = "mnist-index" # 索引名称

# 创建新索引
print(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
print(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
print(f"已成功连接到索引 '{index_name}'。")
# 从 scikit-learn 库中导入 load_digits 函数
# 这个函数用于加载著名的手写数字数据集 MNIST
from sklearn.datasets import load_digits

# 使用 load_digits 函数加载 MNIST 数据集
# n_class=10 表示加载全部 10 个数字类别(0-9)
digits = load_digits(n_class=10)

# 获取数据集中的特征数据
# X 是一个二维数组,每行代表一个样本,每个样本是一个 64 维的向量(8x8 像素展平)
X = digits.data

# 获取数据集中的标签
# y 是一个一维数组,包含每个样本对应的真实数字标签(0-9)
y = digits.target

# 初始化一个空列表,用于存储转换后的向量数据
vectors = []
for i in range(len(X)):
    vector_id = str(i)
    vector_values = X[i].tolist()
    metadata = {"label": int(y[i])}
    # 改为 Pinecone 推荐的字典格式
    vectors.append({
        "id": vector_id,
        "values": vector_values,
        "metadata": metadata
    })

# 定义批处理大小,每批最多包含 1000 个向量
# 这是为了避免一次性向 Pinecone 发送过多数据,可能导致请求超时或失败
batch_size = 1000

# 使用步长为 batch_size 的 range 函数,实现分批处理
for i in range(0, len(vectors), batch_size):
    # 从 vectors 列表中切片获取一批数据
    batch = vectors[i:i + batch_size]
    # 使用 upsert 方法将这批数据上传到 Pinecone 索引中
    # upsert 操作会插入新的向量或更新已存在的向量
    upsert_result = index.upsert(batch, namespace="")
    print(f"Upsert batch {i//batch_size+1}: {upsert_result}")
# 导入必要的库
import matplotlib.pyplot as plt  # 用于绘图
import numpy as np  # 用于数值计算
from collections import Counter  # 用于计数
import time

# 等待索引状态为 "Ready"
while True:
    desc = pinecone.describe_index(index_name)
    # 直接用属性访问
    status = getattr(desc.status, "ready", None)
    print(f"索引状态: {status}")
    if status is True:
        break
    time.sleep(1)

# 等待索引同步
time.sleep(2)  # 等待索引同步

# upsert 必须在索引 ready 之后，只执行一次
for i in range(0, len(vectors), batch_size):
    batch = vectors[i:i + batch_size]
    upsert_result = index.upsert(batch, namespace="")
    print(f"Upsert batch {i//batch_size+1}: {upsert_result}")

# 检查索引数据量
stats = index.describe_index_stats()
print("索引统计信息：", stats)

# 创建一个手写数字 3 的图像
# 使用 numpy 数组表示一个 8x8 的二维图像
# 255 表示白色像素,0 表示黑色像素
digit_3 = np.array(
    [[0, 0, 255, 255, 255, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 255, 255, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 255, 255, 255, 255, 0, 0]]
)

# 将图像像素值从 0-255 的范围缩放到 0-16 的范围
# 这是为了匹配 MNIST 数据集中使用的像素值范围
digit_3_flatten = np.round((digit_3 / 255.0) * 16).astype(int)

# 将二维图像数组展平成一维列表
# 这是因为 Pinecone 要求输入向量是一维的
query_data = digit_3_flatten.ravel().tolist()
print("查询向量（前10维）：", query_data[:10])

# 打印数据集中第一个标签为3的样本，便于对比
idx_3 = np.where(y == 3)[0][0]
print("数据集中第一个'3'的向量（前10维）：", X[idx_3][:10].astype(int).tolist())

# 使用准备好的查询向量在 Pinecone 索引中执行搜索
results = index.query(
    vector=query_data,
    top_k=11,  # 返回距离最近的 11 个结果
    include_metadata=True,  # 同时返回每个向量的元数据(包括标签)
    namespace=""
)
# 从搜索结果中提取每个匹配项的标签
labels = [match['metadata']['label'] for match in results['matches']]
if not labels:
    print("未找到任何匹配结果，请检查索引或查询向量。")
else:
    # 打印每个匹配结果的详细信息
    for match, label in zip(results['matches'], labels):
        print(f"id: {match['id']}, distance: {match['score']}, label: {label}")
    # 使用投票机制确定最终的分类结果
    final_prediction = Counter(labels).most_common(1)[0][0]
    # 使用 matplotlib 显示查询图像和预测结果
    plt.imshow(digit_3, cmap='gray')  # 显示灰度图像
    plt.title(f"Predicted digit: {final_prediction}", size=15)  # 设置标题,显示预测结果
    plt.axis('off')  # 关闭坐标轴
    plt.show()  # 展示图像

# 用数据集中第一个'3'的样本做查询
query_data = X[idx_3].tolist()
print("用真实数据集的'3'做查询，前10维：", query_data[:10])
# 使用准备好的查询向量在 Pinecone 索引中执行搜索
results = index.query(
    vector=query_data,
    top_k=11,  # 返回距离最近的 11 个结果
    include_metadata=True,  # 同时返回每个向量的元数据(包括标签)
    namespace=""
)
# 从搜索结果中提取每个匹配项的标签
labels = [match['metadata']['label'] for match in results['matches']]
if not labels:
    print("未找到任何匹配结果，请检查索引或查询向量。")
else:
    # 打印每个匹配结果的详细信息
    for match, label in zip(results['matches'], labels):
        print(f"id: {match['id']}, distance: {match['score']}, label: {label}")
    # 使用投票机制确定最终的分类结果
    final_prediction = Counter(labels).most_common(1)[0][0]
    # 使用 matplotlib 显示查询图像和预测结果
    plt.imshow(digit_3, cmap='gray')  # 显示灰度图像
    plt.title(f"Predicted digit: {final_prediction}", size=15)  # 设置标题,显示预测结果
    plt.axis('off')  # 关闭坐标轴
    plt.show()  # 展示图像

# 检查索引数据量
stats = index.describe_index_stats()
print("索引统计信息：", stats)
